import os
import numpy as np
from pathlib import Path


def split_adsb_tracks(input_dir, output_dir):
    """
    Split ADS-B data files by segment_id into individual track files.

    Args:
        input_dir (str): Directory containing input npy files
        output_dir (str): Directory to save split track files
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    # Get all npy files in input directory
    input_files = list(Path(input_dir).glob('*.npy'))

    if not input_files:
        print(f"No .npy files found in {input_dir}")
        return

    for file_path in input_files:
        print(f"\nProcessing {file_path.name}...")

        # Load data
        data = np.load(file_path, allow_pickle=True)

        # Check if data is loaded correctly
        if data is None or not isinstance(data, np.ndarray):
            print(f"Warning: Failed to load data from {file_path.name}, skipping...")
            continue

        # Print data shape
        print(f"Data shape: {data.shape}")

        # Assuming data.shape is (n_samples, n_features)
        if data.ndim != 2 or data.shape[1] < 6:
            print(f"Warning: Data shape is unexpected in {file_path.name}, skipping...")
            continue

        # Extract columns by index
        time = data[:, 0].astype(np.float64)         # time column
        lat = data[:, 2]
        lon = data[:, 3]
        geoaltitude = data[:, 13]
        segment_id = data[:, 21]
        icao24 = data[:, 1]

        # If icao24 is bytes type, decode it
        if isinstance(icao24[0], bytes):
            icao24 = np.array([x.decode('utf-8') for x in icao24])

        # Get unique segment IDs
        unique_segments = np.unique(segment_id)
        print(f"Found {len(unique_segments)} unique segments")

        for seg_id in unique_segments:
            # Create mask for current segment
            mask = segment_id == seg_id

            # Get icao24 for this segment
            current_icao24 = icao24[mask][0]  # Assuming icao24 is consistent within a segment

            # Stack the required columns
            output_data = np.column_stack((
                time[mask],
                icao24[mask],
                lat[mask],
                lon[mask],
                geoaltitude[mask],
                segment_id[mask]
            ))

            # Create output filename
            output_filename = f"{current_icao24}_{int(seg_id)}.npy"
            output_path = os.path.join(output_dir, output_filename)

            # Save the segment data
            np.save(output_path, output_data)
            print(f"Saved segment {seg_id} to {output_filename}")

        print(f"Finished processing {file_path.name}")


# Example usage
if __name__ == "__main__":
    input_directory = "D:/pythonProject/粒子滤波/ADS-B/Tset-npy"  # Replace with your input directory
    output_directory = "D:/pythonProject/粒子滤波/ADS-B/拆分-npy"  # Replace with your output directory

    split_adsb_tracks(input_directory, output_directory)